# ============================================================
#  final_gpbo.py   — 4‑层 DP + Gaussian‑Process Bayesian Opt.
#  （包含完整输出：αβ日志、散点图、调度曲线 Excel & PNG）
# ============================================================
"""
使用 GP‑Bayesian Optimization (EI) 代替原 PSO，在极低评估预算内
搜索二维决策 (α, β)。

输出内容（均写入 OUTPUT_DIR）：
    • bo_eval_log.xlsx            —— 每次评估的 α, β, 期望收益
    • bo_revenue_scatter.png      —— 三维散点图
    • dispatch_all_scenarios.xlsx —— α★,β★ 下 25 个场景 7×24 调度曲线
    • scen_*_dispatch.png×5       —— 权重最高的前 5 场景曲线图

依赖： pip install scikit-optimize scikit-learn pandas openpyxl matplotlib joblib
"""

# ============ 用户自行修改 =====================
OUTPUT_DIR = r"./嵌套收益模型"  # ← 结果保存目录
N_CALLS_GP = 40  # GP‑BO 总评估次数 (含初始)
INITIAL_PTS = 10  # Sobol/LH 起始样本数
REPORT_EVERY = 2  # L1 监控打印步长 (h)
SCENE_DEBUG = {(0, 0)}  # 只监控这些场景的 L1 进度
# ===============================================

import json
import math
import os
import time
import warnings
from functools import lru_cache
from itertools import product
from typing import Dict, Tuple

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# ---------- 业务逻辑 & 数据 ----------------------
# 假设 dp_core.py 定义了所有 DP 常量、数据及 instant_revenue()
from dp_core import (
    FLOW,
    GRID,
    HOURS,  # DP 结构
    KG,
    LAYER_RES,
    P_DA,
    P_RT,
    PMAX,  # 基础参数 (np.ndarray)
    PMIN,
    RATIO_MAX,
    SCN_P,  # 场景数据
    WIND,
    IterativeDP,  # 辅助函数和类
    decompose,
    instant_revenue,
)
from joblib import Parallel, delayed, parallel_backend
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401  (启用 3D)
from skopt import gp_minimize
from skopt.space import Real
from skopt.utils import use_named_args


# ============ 监控工具（仅 L1 使用） =============
def _report(layer: str, t: int, k: int, j: int):
    if (k, j) not in SCENE_DEBUG:
        return
    if t % REPORT_EVERY == 0 or t == HOURS:
        print(f"[{layer}] (scen {k},{j}) t = {t:02d}/{HOURS}", flush=True)


# ============ build_solvers — 迭代DP版本 =======
EMPTY_PLAN = np.zeros((len(FLOW), HOURS), dtype=np.float32)  # (7,24)


def build_iterative_solver_for_gpbo(env: Dict, k: int, j: int):
    """为GPBO构建迭代DP求解器"""
    return IterativeDP(env, k, j)


# ============ 单场景评估 (返回收益 & 计划) =========


def _solve_one_scene(k: int, j: int, alpha: float, beta: float):
    water_day = float(np.sum(FLOW))  # 7 库日流量之和
    wind_vec = WIND[j]
    wind_day = float(np.sum(wind_vec))  # 风电日量
    total_day = water_day + wind_day

    con_curve, da_curve = decompose(alpha, beta, total_day)

    env = {
        "con": con_curve,
        "da": da_curve,
        "p_da": P_DA[k],
        "p_rt": P_RT[k],
        "wind": wind_vec,
    }

    solver = build_iterative_solver_for_gpbo(env, k, j)
    best_val, best_plan = solver.solve()

    return SCN_P[k, j] * best_val, best_plan


# ============ GP‑BO 目标函数 (只返回收益) ==========


def cost_one_particle(alpha: float, beta: float) -> float:
    # ① 越界直接给大罚值
    if alpha < 0 or beta < 0 or alpha + beta > 1:
        return 1e12

    # ② 并行跑 25 个场景，返回 [(revenue, plan), ...]
    with parallel_backend("loky", inner_max_num_threads=1):
        scene_vals = Parallel(n_jobs=-1, batch_size=1)(
            delayed(_solve_one_scene)(k, j, alpha, beta)
            for k, j in product(range(5), range(5))
        )

    # ③ 求期望收益；plan 用不到，只保留日志

    total_revenue = sum(rev for rev, _ in scene_vals)

    # ④ 记录 α-β-收益 轨迹（全局 eval_log 在 main 外定义）
    eval_log.append((alpha, beta, total_revenue))

    # ⑤ skopt 以"最小化"为目标 → 返回 –收益
    return -total_revenue


# ============ 搜索空间 + skopt 包装 ================


def _project(u: float, v: float):
    a, b = abs(u), abs(v)
    s = a + b
    if s > 1:
        a, b = a / s, b / s
    return a, b


@use_named_args([Real(0.0, 1.0, name="u"), Real(0.0, 1.0, name="v")])
def objective(u, v):
    a, b = _project(u, v)
    return cost_one_particle(a, b)


# ============ 主流程 =================================

eval_log: list = []  # [(α,β,收益)]  — 在 cost_one_particle 填


def main():
    warnings.filterwarnings("ignore", category=UserWarning)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("=== GP‑Bayesian Optimization 开始 ===")
    t0 = time.time()

    res = gp_minimize(
        objective,
        dimensions=[Real(0.0, 1.0, name="u"), Real(0.0, 1.0, name="v")],
        n_initial_points=INITIAL_PTS,
        n_calls=N_CALLS_GP,
        acq_func="EI",
        random_state=42,
        verbose=True,
    )

    best_alpha, best_beta = _project(*res.x)
    best_revenue = -res.fun

    print(
        f"\n>>> α★,β★ = ({best_alpha:.4f}, {best_beta:.4f})  收益={best_revenue:,.2f}"
    )
    print(f"耗时 {(time.time() - t0):.1f}s  评估 {len(res.func_vals)} 次")

    # ---------- 保存 αβ 日志 --------------
    df_log = pd.DataFrame(eval_log, columns=["alpha", "beta", "revenue"])
    df_log.to_excel(os.path.join(OUTPUT_DIR, "bo_eval_log.xlsx"), index=False)

    # ---------- 绘制散点图 ----------------
    fig = plt.figure(figsize=(6, 5))
    ax = fig.add_subplot(111, projection="3d")
    ax.scatter(df_log.alpha, df_log.beta, df_log.revenue, s=25)
    ax.set_xlabel("α"), ax.set_ylabel("β"), ax.set_zlabel("Revenue")
    plt.title("GP‑BO Search Trajectory")
    plt.tight_layout()
    plt.savefig(os.path.join(OUTPUT_DIR, "bo_revenue_scatter.png"), dpi=300)
    plt.close(fig)

    # ---------- 重新计算 25 场景调度曲线（主进程） ----------
    dispatch_dict = {}
    for k, j in product(range(5), range(5)):
        _, plan = _solve_one_scene(k, j, best_alpha, best_beta)
        dispatch_dict[(k, j)] = plan

    # ---------- 写 Excel (所有场景) ----------
    xls_path = os.path.join(OUTPUT_DIR, "dispatch_all_scenarios.xlsx")
    with pd.ExcelWriter(xls_path, engine="openpyxl") as writer:
        for (k, j), plan in dispatch_dict.items():
            df_plan = pd.DataFrame(
                plan.T,
                columns=[f"Res{r + 1}" for r in range(len(FLOW))],
                index=[f"H{h + 1}" for h in range(HOURS)],
            )
            df_plan.to_excel(writer, sheet_name=f"scen_{k}_{j}")

    # ---------- 取权重最高前 5 场景画图 ----------
    flat_prob = sorted(
        [(SCN_P[k, j], k, j) for k in range(5) for j in range(5)], reverse=True
    )[:5]
    for p, k, j in flat_prob:
        plan = dispatch_dict[(k, j)]
        plt.figure(figsize=(7, 4))
        for r in range(len(FLOW)):
            plt.plot(range(1, HOURS + 1), plan[r], label=f"Res{r + 1}")
        plt.title(f"Dispatch  scen({k},{j})  P={p:.3f}")
        plt.xlabel("Hour"), plt.ylabel("Relative Q")
        plt.legend(ncol=2, fontsize=8)
        plt.tight_layout()
        plt.savefig(os.path.join(OUTPUT_DIR, f"scen_{k}_{j}_dispatch.png"), dpi=300)
        plt.close()

    print("所有输出已写入", OUTPUT_DIR)


# ===========================
if __name__ == "__main__":
    from multiprocessing import freeze_support

    freeze_support()
    main()